import pandas as pd
import argparse
import json
import os
import sys
from tkinter import Tk, filedialog
from tkinter.messagebox import showerror

print("使用前务必删除冗余表头：\n- 手动操作：打开Excel → 定位包含【四平市普通高中xxxx.xx.xx某某考试联榜成绩】的行 → 右键删除 → 下方单元格上移")
print("中宏软件技术开源社区-Zeta开发")
# 配置文件默认路径
DEFAULT_CONFIG = "config.json"

def load_config(config_path):
    """加载配置文件"""
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        
        # 验证必要字段
        required_sections = ['base_columns', 'subjects', 'metrics', 'column_mapping']
        for section in required_sections:
            if section not in config:
                raise ValueError(f"配置文件中缺少必要字段: {section}")
        
        return config
    except Exception as e:
        print(f"加载配置文件失败: {str(e)}")
        sys.exit(1)

def read_data(input_path, header_levels):
    """读取Excel数据"""
    try:
        return pd.read_excel(input_path, header=header_levels)
    except Exception as e:
        print(f"读取文件失败: {str(e)}")
        sys.exit(1)

def build_filter_condition(df, config, filters):
    """构建过滤条件"""
    conditions = []
    
    # 学校过滤
    if 'school' in filters and filters['school']:
        school_col = tuple(config['column_mapping']['学校'])
        if school_col not in df.columns:
            raise KeyError(f"列 {school_col} 不存在")
        conditions.append(df[school_col] == filters['school'])
    
    # 班级过滤
    if 'class_' in filters and filters['class_']:
        if 'school' not in filters or not filters['school']:
            raise ValueError("班级过滤需要同时指定学校")
        class_col = tuple(config['column_mapping']['班级'])
        if class_col not in df.columns:
            raise KeyError(f"列 {class_col} 不存在")
        conditions.append(df[class_col] == filters['class_'])
    
    # 选科组合过滤
    if 'combination' in filters and filters['combination']:
        combo_col = tuple(config['column_mapping']['选科组合'])
        if combo_col not in df.columns:
            raise KeyError(f"列 {combo_col} 不存在")
        conditions.append(df[combo_col] == filters['combination'])
    
    if not conditions:
        return pd.Series([True]*len(df), index=df.index)
    
    return pd.concat(conditions, axis=1).all(axis=1)

def select_columns(df, config):
    """选择需要保留的列"""
    selected = []
    
    # 添加基础列
    for col_name in config['base_columns']:
        col_path = config['column_mapping'].get(col_name)
        if not col_path:
            continue
        col_tuple = tuple(col_path)
        if col_tuple in df.columns:
            selected.append(col_tuple)
    
    # 添加科目指标组合
    for subject in config['subjects']:
        for metric in config['metrics']:
            col_tuple = (subject, metric)
            if col_tuple in df.columns:
                selected.append(col_tuple)
    
    return df[selected]

def save_file_interactive(default_name):
    """交互模式保存文件"""
    root = Tk()
    root.withdraw()
    try:
        file_path = filedialog.asksaveasfilename(
            defaultextension='.xlsx',
            filetypes=[('Excel Files', '*.xlsx')],
            initialfile=default_name
        )
    except Exception as e:
        showerror("错误", f"文件保存失败: {str(e)}")
        return None
    finally:
        root.destroy()
    return file_path

def handle_command_line(args, config):
    """处理命令行模式"""
    # 读取数据
    df = read_data(args.input, [0, 1])
    
    # 构建过滤条件
    try:
        filters = {
            'school': args.school,
            'class_': args.class_,
            'combination': args.combination
        }
        condition = build_filter_condition(df, config, filters)
    except Exception as e:
        print(f"过滤条件错误: {str(e)}")
        sys.exit(1)
    
    # 应用过滤
    filtered_df = df[condition]
    if filtered_df.empty:
        print("警告: 没有匹配的数据")
    
    # 选择列
    result_df = select_columns(filtered_df, config)
    
    # 处理列名
    result_df.columns = ['_'.join(map(str, col)).strip('_') for col in result_df.columns.values]
    
    # 保存结果
    if args.output:
        output_path = args.output
    else:
        output_path = "输出结果.xlsx"
    
    try:
        result_df.to_excel(output_path, index=False)
        print(f"结果已保存到: {os.path.abspath(output_path)}")
    except Exception as e:
        print(f"保存文件失败: {str(e)}")
        sys.exit(1)

def interactive_mode(config):
    """交互模式处理"""
    # 获取输入文件
    while True:
        input_path = input("请输入源文件路径（留空使用文件对话框）: ").strip()
        if not input_path:
            root = Tk()
            root.withdraw()
            input_path = filedialog.askopenfilename(
                filetypes=[('Excel Files', '*.xlsx')]
            )
            root.destroy()
            if not input_path:
                print("必须指定输入文件！")
                continue
        
        if os.path.exists(input_path):
            break
        print("文件不存在，请重新输入！")
    
    # 读取数据
    try:
        df = read_data(input_path, [0, 1])
    except Exception as e:
        print(str(e))
        sys.exit(1)
    
    # 选择过滤模式
    while True:
        print("\n请选择过滤条件类型:")
        print("1. 单一学校")
        print("2. 单一学校+班级")
        print("3. 单一学校+选科组合")
        print("4. 总选科组合（所有学校）")
        choice = input("请输入选项 (1-4): ").strip()
        
        if choice in ['1', '2', '3', '4']:
            break
        print("无效输入，请重新选择！")
    
    filters = {}
    # 获取过滤条件
    try:
        if choice in ['1', '2', '3']:
            school = input("请输入学校名称: ").strip()
            if not school:
                raise ValueError("学校名称不能为空")
            filters['school'] = school
            
            if choice == '2':
                class_ = input("请输入班级名称: ").strip()
                if not class_:
                    raise ValueError("班级名称不能为空")
                filters['class_'] = class_
            
            if choice == '3':
                combination = input("请输入选科组合: ").strip()
                if not combination:
                    raise ValueError("选科组合不能为空")
                filters['combination'] = combination
        
        if choice == '4':
            combination = input("请输入选科组合: ").strip()
            if not combination:
                raise ValueError("选科组合不能为空")
            filters['combination'] = combination
    except Exception as e:
        print(f"输入错误: {str(e)}")
        sys.exit(1)
    
    # 构建过滤条件
    try:
        condition = build_filter_condition(df, config, filters)
    except Exception as e:
        print(f"过滤条件错误: {str(e)}")
        sys.exit(1)
    
    # 应用过滤
    filtered_df = df[condition]
    if filtered_df.empty:
        print("警告: 没有匹配的数据")
    
    # 选择列
    result_df = select_columns(filtered_df, config)
    
    # 处理列名
    print("开始进行数据清洗")
    result_df.columns = ['_'.join(map(str, col)).strip('_') for col in result_df.columns.values]
    print("数据清洗完成")
    
    # 保存文件
    default_name = "筛选结果.xlsx"
    output_path = save_file_interactive(default_name)
    while not output_path:
        print("必须指定保存路径！")
        output_path = save_file_interactive(default_name)
    
    try:
        result_df.to_excel(output_path, index=False)
        print(f"结果已保存到: {os.path.abspath(output_path)}")
    except Exception as e:
        print(f"保存文件失败: {str(e)}")
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(
        description="四平市普通高中联考榜单数据分离软件",
        formatter_class=argparse.RawTextHelpFormatter,
        epilog="使用前务必删除冗余表头：\n- 手动操作：打开Excel → 定位包含【四平市普通高中xxxx.xx.xx某某考试联榜成绩】的行 → 右键删除 → 下方单元格上移\n"
        "示例:\n"
        "如是预编译版本请不要按照下面提示内的python main.py 输入，直接运行main.exe+参数即可\n"
        "  命令行模式: \n"
        "    python main.py --input 输入数据.xlsx --school '四平市第三高级中学' --combination '物化政'\n"
        "    python main.py --input data.xlsx --class 3班 --school 四平市实验中学\n\n"
        "  交互模式: \n"
        "    python main.py\n"
        "中宏软件技术开源社区-Zeta开发"
    )
    parser.add_argument('--input', help='输入文件路径')
    parser.add_argument('--output', help='输出文件路径')
    parser.add_argument('--school', help='筛选学校名称')
    parser.add_argument('--class', dest='class_', help='筛选班级名称（需同时指定--school）')
    parser.add_argument('--combination', help='筛选选科组合名称')
    parser.add_argument('--config', default=DEFAULT_CONFIG, help=f'配置文件路径（默认: {DEFAULT_CONFIG}）')
    args = parser.parse_args()
    
    # 加载配置
    config = load_config(args.config)
    
    # 模式判断
    if args.input or any([args.school, args.class_, args.combination]):
        handle_command_line(args, config)
    else:
        interactive_mode(config)

if __name__ == "__main__":
    main()